MODULE sefmm_cheb_AO

   USE fmm_global_paras
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: sefmm_get_AO_Tn_ints

   REAL(REALK), ALLOCATABLE :: ECoefX(:,:,:,:) 
   REAL(REALK), ALLOCATABLE :: ECoefY(:,:,:,:)
   REAL(REALK), ALLOCATABLE :: ECoefZ(:,:,:,:)

   REAL(REALK), ALLOCATABLE :: MIntX(:,:,:)
   REAL(REALK), ALLOCATABLE :: MIntY(:,:,:)
   REAL(REALK), ALLOCATABLE :: MIntZ(:,:,:)

   REAL(REALK), ALLOCATABLE :: SintX(:,:,:,:)
   REAL(REALK), ALLOCATABLE :: SintY(:,:,:,:)
   REAL(REALK), ALLOCATABLE :: SintZ(:,:,:,:)

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_init_AO_Tn(basis,npts)

      IMPLICIT NONE
      TYPE(fmm_basis), INTENT(IN) :: basis
      INTEGER(INTK),   INTENT(IN) :: npts

      INTEGER(INTK) :: ndim, MaxAngl, MaxSgm2

      MaxAngl = basis%maxangl 
      MaxSgm2 = basis%maxsgm2 

      ALLOCATE(ECoefX(0:MaxAngl, 0:MaxAngl, 0:MaxAngl*2, MaxSgm2))
      ALLOCATE(ECoefY(0:MaxAngl, 0:MaxAngl, 0:MaxAngl*2, MaxSgm2))
      ALLOCATE(ECoefZ(0:MaxAngl, 0:MaxAngl, 0:MaxAngl*2, MaxSgm2))

      ndim = 2*MaxAngl
      ALLOCATE(MIntX(0:ndim, MaxSgm2, npts))
      ALLOCATE(MIntY(0:ndim, MaxSgm2, npts))
      ALLOCATE(MIntZ(0:ndim, MaxSgm2, npts))

      ALLOCATE(SintX(0:MaxAngl, 0:MaxAngl, MaxSgm2, npts))
      ALLOCATE(SintY(0:MaxAngl, 0:MaxAngl, MaxSgm2, npts))
      ALLOCATE(SintZ(0:MaxAngl, 0:MaxAngl, MaxSgm2, npts))

   END SUBROUTINE sefmm_init_AO_Tn

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_free_AO_Tn

      IMPLICIT NONE
      IF (ALLOCATED(ECoefX)) DEALLOCATE(ECoefX)
      IF (ALLOCATED(ECoefY)) DEALLOCATE(ECoefY)
      IF (ALLOCATED(ECoefZ)) DEALLOCATE(ECoefZ)

      IF (ALLOCATED(MIntX)) DEALLOCATE(MIntX)
      IF (ALLOCATED(MIntY)) DEALLOCATE(MIntY)
      IF (ALLOCATED(MIntZ)) DEALLOCATE(MIntZ)

      IF (ALLOCATED(SintX)) DEALLOCATE(SintX)
      IF (ALLOCATED(SintY)) DEALLOCATE(SintY)
      IF (ALLOCATED(SintZ)) DEALLOCATE(SintZ)

   END SUBROUTINE sefmm_free_AO_Tn

!-------------------------------------------------------------------------------

   SUBROUTINE sefmm_get_AO_Tn_ints(basis)

      USE fmm_integral_utils, ONLY: fmm_get_prim_batch,       &
                                    fmm_build_Ecoef1,         &
                                    fmm_build_Ecoef2
      USE fefmm_box_utils,    ONLY: fefmm_get_boxes,          &
                                    fefmm_free_boxes,         &
                                    fefmm_sh_pair_in_box

      IMPLICIT NONE
      TYPE(scheme_paras), INTENT(IN) :: scheme
      TYPE(fmm_basis),    INTENT(IN) :: basis
      TYPE(fmm_sh_pairs), INTENT(IN) :: sh_pairs(:)



      CALL sefmm_init_AO_Tn(basis,npts)



      DO ij = 1, SIZE(sh_pairs)

         IF (.NOT.fefmm_sh_pair_in_box(sh_pairs(ij),boxes(:,i),grain)) CYCLE

         Ish = sh_pairs(ij)%I
         Jsh = sh_pairs(ij)%J





      END DO


      CALL sefmm_free_AO_Tn(basis,npts)


!-------------------------------------------------------------------------------

!fixme
      CHARACTER(LEN=13), PARAMETER :: Name = 'FE_AO_overlap'

      REAL(REALK), POINTER :: boxes(:,:)
      TYPE(fmm_prim_batch) :: batch(basis%MaxSgm2)
      CHARACTER(LEN=255) :: FBuf
      INTEGER(INTK) :: Ish, Jsh, NPrim, IAnglA, IAnglB, FEdim, LIPN
      INTEGER(INTK) :: i, ij, k
      REAL(REALK)   :: grain, FEgrain

print *, 'Building raw AO_FE overlap integrals...'

      grain = scheme%grain
      FEdim = scheme%FEdim
      LIPN = scheme%lipn
      FEgrain = scheme%grain / (FEdim -1)

      FBuf = TRIM(Name)//".fefmm"
      OPEN(UNIT=LUINTM, FILE=TRIM(FBuf), STATUS='REPLACE',  &
           ACCESS='SEQUENTIAL', FORM='UNFORMATTED')
      REWIND(LUINTM)

      CALL fefmm_get_boxes(scheme%grain,sh_pairs,boxes)

      DO i = 1, SIZE(boxes,2)
!print *, 'box:', i, 'of', size(boxes,2)
!print *, 'box:', boxes(:,i)

         DO ij = 1, SIZE(sh_pairs)
!print *, 'shell pair:', ij, 'of', size(sh_pairs)

            IF (.NOT.fefmm_sh_pair_in_box(sh_pairs(ij),boxes(:,i),grain)) CYCLE

            Ish = sh_pairs(ij)%I
            Jsh = sh_pairs(ij)%J
  
            ! Prepare batch of primitives
            CALL fmm_get_prim_batch(basis,Ish,Jsh,batch,NPrim)

            ! Build Hermite expansion E-coefficients
            IAnglA = basis%KType(Ish)
            IAnglB = basis%KType(Jsh)
            IF (IAnglA >= IAnglB) THEN
               CALL fmm_build_Ecoef1(batch,NPrim,IAnglA,IAnglB,   &
                                     ECoefX,ECoefY,ECoefZ)
            ELSE
               CALL fmm_build_Ecoef2(batch,NPrim,IAnglA,IAnglB,   &
                                     ECoefX,ECoefY,ECoefZ)
            END IF
  
            ! Build 1D Hermite GTO/FE overlap integrals
            DO k = 1, NPrim
               batch(k)%PC(:) = batch(k)%P(:) - (boxes(:,i)-grain*half)
            END DO
           
            CALL fmm_build_AO_FE_Hermite(batch,NPrim,IAnglA+IAnglB,   &
                                         FEdim,LIPN,FEgrain)

            CALL fmm_build_AO_FE_1D_components(IAnglA,IAnglB,NPrim,FEdim)
            CALL fmm_store_AO_FE_1D_components(IAnglA,IAnglB,NPrim,FEdim)

         END DO
      END DO

      CLOSE(UNIT=LUINTM, STATUS='KEEP')
      NULLIFY (boxes)
      CALL fefmm_free_boxes

   END SUBROUTINE fefmm_build_AO_FE_overlaps

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_contract_FEpot(scheme,basis,sh_pairs,box,V,Jmat)

      USE fmm_integral_utils, ONLY: fmm_get_prim_batch
      USE fefmm_box_utils,    ONLY: fefmm_sh_pair_in_box

      IMPLICIT NONE
      TYPE(scheme_paras), INTENT(IN)    :: scheme
      TYPE(fmm_basis),    INTENT(IN)    :: basis
      TYPE(fmm_sh_pairs), INTENT(IN)    :: sh_pairs(:)
      REAL(REALK),        INTENT(IN)    :: box(3)
      REAL(REALK),        INTENT(IN)    :: V(:)
      REAL(REALK),        INTENT(INOUT) :: Jmat(:,:)

      TYPE(fmm_prim_batch) :: batch(basis%MaxSgm2)
      INTEGER(INTK) :: Ish, Jsh, NPrim, IAnglA, IAnglB, FEdim
      INTEGER(INTK) :: i, ij, k
      REAL(REALK)   :: grain

      grain = scheme%grain
      FEdim = scheme%FEdim

      DO ij = 1, SIZE(sh_pairs)
!print *, 'shell pair:', ij, 'of', size(sh_pairs)

         IF (.NOT. fefmm_sh_pair_in_box(sh_pairs(ij),box,grain)) CYCLE

         Ish = sh_pairs(ij)%I
         Jsh = sh_pairs(ij)%J
         CALL fmm_get_prim_batch(basis,Ish,Jsh,batch,NPrim)

         IAnglA = basis%KType(Ish)
         IAnglB = basis%KType(Jsh)
         CALL fmm_read_AO_FE_1D_components(IAnglA,IAnglB,NPrim,FEdim)
         CALL fefmm_update_J_matrix_exact(basis,batch,Ish,Jsh,NPrim,FEdim,V,Jmat)

      END DO

   END SUBROUTINE fefmm_contract_FEpot

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_build_AO_FE_1D_components(IAnglA,IAnglB,NPrim,FEdim)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: IAnglA, IAnglB, NPrim, FEdim

      INTEGER(INTK) :: It1, It2, It, IJ, k

      SintX(:,:,:,:) = zero
      SintY(:,:,:,:) = zero
      SintZ(:,:,:,:) = zero

      DO It1 = 0, IAnglA ! s,p,d,f,,,
         DO It2 = 0, IAnglB
            DO It = 0, It1+It2
               DO IJ = 1, NPrim
                  DO k = 1, FEdim

                     SintX(It1,It2,IJ,k) = SintX(It1,It2,IJ,k)   +    &
                                           ECoefX(It1,It2,It,IJ) *    &
                                           MIntX(It,IJ,k)

                     SintY(It1,It2,IJ,k) = SintY(It1,It2,IJ,k)   +    &
                                           ECoefY(It1,It2,It,IJ) *    &
                                           MIntY(It,IJ,k)
   
                     SintZ(It1,It2,IJ,k) = SintZ(It1,It2,IJ,k)   +    & 
                                           ECoefZ(It1,It2,It,IJ) *    &
                                           MIntZ(It,IJ,k)
                  END DO
               END DO
            END DO
         END DO
      END DO

   END SUBROUTINE fmm_build_AO_FE_1D_components

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_store_AO_FE_1D_components(IAnglA,IAnglB,NPrim,FEdim)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: IAnglA, IAnglB, NPrim, FEdim

      INTEGER(INTK) :: It1, It2, IJ, k

      WRITE(LUINTM) IAnglA, IAnglB, NPrim, FEdim

      DO It1 = 0, IAnglA
         DO It2 = 0, IAnglB
            DO IJ = 1, NPrim
               DO k = 1, FEdim

                  WRITE(LUINTM) SintX(It1,It2,IJ,k),  &
                                SintY(It1,It2,IJ,k),  &
                                SintZ(It1,It2,IJ,k)

!         print *, it1, it2, ij, k
!         print *, SintX(It1,It2,IJ,k),  &
!                  SintY(It1,It2,IJ,k),  &
!                  SintZ(It1,It2,IJ,k)

               END DO
            END DO
         END DO
      END DO

   END SUBROUTINE fmm_store_AO_FE_1D_components

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_read_AO_FE_1D_components(IAnglA,IAnglB,NPrim,FEdim)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN) :: IAnglA, IAnglB, NPrim, FEdim

      INTEGER(INTK) :: It1, It2, IJ, k
      INTEGER(INTK) :: IAnglA0, IAnglB0, NPrim0, FEdim0

      READ(LUINTM) IAnglA0, IAnglB0, NPrim0, FEdim0
      IF (IAnglA0 /= IAnglA) CALL fmm_quit('IO error in AO_FE')
      IF (IAnglB0 /= IAnglB) CALL fmm_quit('IO error in AO_FE')
      IF (NPrim0 /= NPrim) CALL fmm_quit('IO error in AO_FE')
      IF (FEdim0 /= FEdim) CALL fmm_quit('IO error in AO_FE')

      DO It1 = 0, IAnglA
         DO It2 = 0, IAnglB
            DO IJ = 1, NPrim
               DO k = 1, FEdim

                  READ(LUINTM) SintX(It1,It2,IJ,k),  &
                               SintY(It1,It2,IJ,k),  &
                               SintZ(It1,It2,IJ,k)

               END DO
            END DO
         END DO
      END DO

   END SUBROUTINE fmm_read_AO_FE_1D_components

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_update_J_matrix_exact(basis,batch,Ish,Jsh,NPrim,FEdim,V,Jmat)

      IMPLICIT NONE
      TYPE(fmm_basis),      INTENT(IN)    :: basis
      TYPE(fmm_prim_batch), INTENT(IN)    :: batch(:)
      INTEGER(INTK),        INTENT(IN)    :: Ish, Jsh, NPrim, FEdim
      REAL(REALK),          INTENT(IN)    :: V(:)
      REAL(REALK),          INTENT(INOUT) :: Jmat(:,:)

      REAL(REALK) :: tmp(FEdim**3)
      INTEGER :: IAnglA, IAnglB
      INTEGER :: IL1, IL2, IL3
      INTEGER :: IeX, IeY, IeZ
      INTEGER :: It1, Iu1, Iv1
      INTEGER :: It2, Iu2, Iv2
      INTEGER :: IJ, kx, ky, kz, kxyz
      INTEGER :: IL2Temp
      INTEGER :: p, q
      LOGICAL :: IEqJ

      IAnglA = basis%KType(Ish)
      IAnglB = basis%KType(Jsh)

      IEqJ = (ISh == JSh)
      p = basis%KLoc_Car(ISh)

      DO IL1 = basis%LtuvMin_Car(IAnglA), basis%LtuvMax_Car(IAnglA)
         p = p + 1
         q = basis%KLoc_Car(JSh)
 
         It1 = basis%Lt(IL1)
         Iu1 = basis%Lu(IL1)
         Iv1 = basis%Lv(IL1)

         IL2Temp = basis%LtuvMax_Car(IAnglB)
         IF (IEqJ) IL2Temp = IL1
         DO IL2 = basis%LtuvMin_Car(IAnglB), IL2Temp
            q = q + 1
            It2 = basis%Lt(IL2)
            Iu2 = basis%Lu(IL2)
            Iv2 = basis%Lv(IL2)

            tmp(:) = zero
            primitives: DO IJ = 1, NPrim
               kxyz = 0
               DO kx = 1, FEdim
                  DO ky = 1, FEdim
                     DO kz = 1, FEdim
                        kxyz = kxyz +1
                        tmp(kxyz) = tmp(kxyz) + SintX(It1,It2,IJ,kx) *    &
                                                SintY(Iu1,Iu2,IJ,ky) *    &
                                                SintZ(Iv1,Iv2,IJ,kz) *    &
                                                batch(IJ)%CCoefAB
                     END DO
                  END DO
               END DO
            END DO primitives

!print *, Ish, Jsh, p,q, DOT_PRODUCT(tmp(:),V(:))
            Jmat(p,q) = Jmat(p,q) + DOT_PRODUCT(tmp(:),V(:))
            Jmat(q,p) = Jmat(p,q)

         END DO
      END DO

   END SUBROUTINE fefmm_update_J_matrix_exact

!-------------------------------------------------------------------------------

   SUBROUTINE fmm_build_AO_FE_Hermite(batch,NPrim,MaxAng,FEdim,LIPN,lambda)

      IMPLICIT NONE
      TYPE(fmm_prim_batch), INTENT(IN) :: batch(:)
      INTEGER(INTK),        INTENT(IN) :: NPrim, MaxAng, FEdim, LIPN
      REAL(REALK),          INTENT(IN) :: lambda

      SELECT CASE (LIPN)
         CASE (1)
            CALL fefmm_LIP1_AO_FE(batch,NPrim,MaxAng,FEdim,lambda)
         CASE (2)
            CALL fefmm_LIP2_AO_FE(batch,NPrim,MaxAng,FEdim,lambda)
         CASE (3)
            CALL fefmm_LIP3_AO_FE(batch,NPrim,MaxAng,FEdim,lambda)
         CASE DEFAULT
            CALL fmm_quit('NLIP too large in AO_FE!')
      END SELECT

   END SUBROUTINE fmm_build_AO_FE_Hermite

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_LIP1_AO_FE(batch,NPrim,MaxAng,FEdim,lambda)

      USE fefmm_LIP1_hermite, ONLY: LIP1_M10, LIP1_M11, LIP1_M12,  &
                                    LIP1_M13, LIP1_M14

      IMPLICIT NONE
      TYPE(fmm_prim_batch), INTENT(IN) :: batch(:)
      INTEGER(INTK),        INTENT(IN) :: NPrim, MaxAng, FEdim
      REAL(REALK),          INTENT(IN) :: lambda

      INTEGER(INTK) :: ij, k, xyz, t
      REAL(REALK)   :: Q(3), tmp1(3), tmp2(3), p

      MIntX(:,:,:) = Zero
      MIntY(:,:,:) = Zero
      MIntZ(:,:,:) = Zero

      IF (MaxAng > 2) CALL fmm_quit('more than <FE|pp> NYI!')

      DO t = 0, MaxAng
         DO ij = 1, NPrim
            p = batch(ij)%ExpntP
            DO k = 1, FEdim
               tmp1(:) = zero
               tmp2(:) = zero
               DO xyz = 1, 3
                  IF (k < FEdim) THEN
                     ! LIP function 1
                     Q(xyz) = batch(ij)%PC(xyz) - (k-1)*lambda  ! P - P1
!print *, 'LIP1:', xyz, p, lambda, Q(xyz)
                     IF (t == 0) THEN
                        tmp1(xyz) = LIP1_M10(p,Q(xyz),lambda)
                     ELSE IF (t == 1) THEN
                        tmp1(xyz) = LIP1_M11(p,Q(xyz),lambda)
                     ELSE IF (t == 2) THEN
                        tmp1(xyz) = LIP1_M12(p,Q(xyz),lambda)
                     ELSE IF (t == 3) THEN
                        tmp1(xyz) = LIP1_M13(p,Q(xyz),lambda)
                     ELSE IF (t == 4) THEN
                        tmp1(xyz) = LIP1_M14(p,Q(xyz),lambda)
                     END IF
                  END IF
                  IF (k > 1) THEN
                     ! LIP function 2
                     Q(xyz) = batch(ij)%PC(xyz) - (k-1)*lambda  ! P - P2
!print *, 'LIP2:', xyz, p, lambda, Q(xyz)
                     IF (t == 0) THEN
                        tmp2(xyz) = LIP1_M10(p,-Q(xyz),lambda)
                     ELSE IF (t == 1) THEN
                        tmp2(xyz) = LIP1_M11(p,-Q(xyz),lambda)
                     ELSE IF (t == 2) THEN
                        tmp2(xyz) = LIP1_M12(p,-Q(xyz),lambda)
                     ELSE IF (t == 3) THEN
                        tmp2(xyz) = LIP1_M13(p,-Q(xyz),lambda)
                     ELSE IF (t == 4) THEN
                        tmp2(xyz) = LIP1_M14(p,-Q(xyz),lambda)
                     END IF
                  END IF
               END DO
               ! Total
               MIntX(t,ij,k) = tmp1(1) + tmp2(1)
               MIntY(t,ij,k) = tmp1(2) + tmp2(2)
               MIntZ(t,ij,k) = tmp1(3) + tmp2(3)
            END DO
         END DO
      END DO
           
   END SUBROUTINE fefmm_LIP1_AO_FE

!-------------------------------------------------------------------------------

   SUBROUTINE fefmm_LIP2_AO_FE(batch,NPrim,MaxAng,FEdim,lambda)

      USE fefmm_LIP2_hermite, ONLY: LIP2_M1, LIP2_M2, LIP2_M3

      IMPLICIT NONE
      TYPE(fmm_prim_batch), INTENT(IN) :: batch(:)
      INTEGER(INTK),        INTENT(IN) :: NPrim, MaxAng, FEdim
      REAL(REALK),          INTENT(IN) :: lambda

      INTEGER(INTK), PARAMETER :: LIPN = 2
      INTEGER(INTK) :: ij, i,j,k, xyz, t
      REAL(REALK)   :: tmp(3,LIPN+1)
      REAL(REALK)   :: Q(3), p

      MIntX(:,:,:) = Zero
      MIntY(:,:,:) = Zero
      MIntZ(:,:,:) = Zero

      DO t = 0, MaxAng
         DO ij = 1, NPrim
            p = batch(ij)%ExpntP
            DO i = 1, (FEdim-1)/LIPN   ! finite elements

               DO xyz = 1, 3
                  Q(xyz) = batch(ij)%PC(xyz) - two*(i-1)*lambda   ! P - P1
                  tmp(xyz,1) = LIP2_M1(t,p,Q(xyz),lambda)
                  tmp(xyz,2) = LIP2_M2(t,p,Q(xyz),lambda)
                  tmp(xyz,3) = LIP2_M3(t,p,Q(xyz),lambda)
               END DO

               DO j = 1, (LIPN+1)
                  k = (i-1)*2 + j 
                  MIntX(t,ij,k) = MIntX(t,ij,k) + tmp(1,j)
                  MIntY(t,ij,k) = MIntY(t,ij,k) + tmp(2,j)
                  MIntZ(t,ij,k) = MIntZ(t,ij,k) + tmp(3,j)
               END DO

            END DO
         END DO
      END DO
           
   END SUBROUTINE fefmm_LIP2_AO_FE

!-------------------------------------------------------------------------------

END MODULE sefmm_cheb_AO
